namespace Nintest.Collections
{
    using System;
    using System.Collections;
    using System.Collections.Generic;
    using System.Linq;

    public class Lookup<TKey, TElement, TGrouping> : IMutableLookup<TKey, TElement, TGrouping>
        where TGrouping : class, IGrouping<TKey, TElement>
    {
        private readonly IDictionary<TKey, TGrouping> _dictionary;
        private readonly IEqualityComparer<TKey> _comparer;
        private TGrouping _defaultKeyGrouping;

        public Lookup()
            : this(EqualityComparer<TKey>.Default)
        {            
        }

        private Lookup(IEqualityComparer<TKey> comparer)
        {
            _comparer = comparer;
            _dictionary = new Dictionary<TKey, TGrouping>(_comparer);
        }

        public int Count
        {
            get
            {
                var count = _dictionary.Count;
                return _defaultKeyGrouping != null ? count + 1 : count;
            }
        }

        public IEnumerable<TElement> this[TKey key]
        {
            get
            {
                TGrouping grouping;
                return TryGetGrouping(key, out grouping) ? grouping : Enumerable.Empty<TElement>();
            }
        }

        public bool TryGetGrouping(TKey key, out TGrouping grouping)
        {
            if (_comparer.Equals(key, default(TKey)))
            {
                grouping = _defaultKeyGrouping;
                return _defaultKeyGrouping != null;
            }

            return _dictionary.TryGetValue(key, out grouping);
        }

        public bool Contains(TKey key)
        {
            return _comparer.Equals(key, default(TKey)) ? _defaultKeyGrouping != null : _dictionary.ContainsKey(key);
        }

        public IEnumerator<IGrouping<TKey, TElement>> GetEnumerator()
        {
            foreach (var kvp in _dictionary)
                yield return kvp.Value;

            if (_defaultKeyGrouping != null)
                yield return _defaultKeyGrouping;
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return ((IEnumerable<IGrouping<TKey, TElement>>)this).GetEnumerator();
        }

        public void Add(TGrouping grouping)
        {
            if (_comparer.Equals(grouping.Key, default(TKey)))
            {
                if (_defaultKeyGrouping != null)
                    ThrowAlreadyContains();
                _defaultKeyGrouping = grouping;
            }
            else
            {
                if (_dictionary.ContainsKey(grouping.Key))
                    ThrowAlreadyContains();
                _dictionary.Add(grouping.Key, grouping);
            }
        }

        private static void ThrowAlreadyContains()
        {
            throw new InvalidOperationException("Lookup already contains a grouping with the given key.");
        }
    }
}